#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Jul 26 20:36:12 2018
@project: 天池比赛-A股主板上市公司公告信息抽取
@group: MZH_314
@author: LHQ
"""
import os

import pandas as pd
import numpy as np

from scipy.sparse import hstack
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import KFold, cross_val_score
from sklearn.feature_selection import SelectKBest
from sklearn.feature_selection import chi2
from sklearn import svm
from sklearn.externals import joblib

from reportIE.utils import seg
from reportIE.feature import Feature


def process_field_docs(docs):
    """分词，方便后续sklearn.Countvertorizer使用
    """
    new_docs = [" ".join(seg(doc)) for doc in docs]
    return new_docs


class FieldFeature:
    """字段特征抽取类
    包括上下文的特定模式特征(通过人为编写的正则表达式定义)，实体
    本身类型特征(日期、数值、百分比之类的)，词语特征(左右邻)

    Attribute
    ---------
    left_feat : sklearn.CountVectorizer 实例
        抽取实体左邻的词语特征
    
    right_feat : sklearn.CountVectorizer 实例
        抽取实体右邻的词语特征

    feat_select : sklearn.feature_select.SelectKBest 实例
        对左右邻词语特征进行特诊选择

    pattern_feat : reportIE.feature.Feature 实例
        抽取实体左右句子的特定模式特征
    """

    def __init__(self, left_feat, right_feat, feat_select):
        self.left_feat = left_feat
        self.right_feat = right_feat
        self.feat_select = feat_select
        
        self.pattern_feat = self.init_pattern_feat()
        
    def init_pattern_feat(self):
        """定义一些模式
        """
        feat = Feature()
        feat.add_pattern("减持")
        feat.add_pattern("增持")
        feat.add_pattern("通过.*?[减增]持")
        feat.add_pattern("\d+股")
        feat.add_pattern("\d+-\d+-\d+")
        feat.add_pattern("后")
        feat.add_pattern("占")
        feat.add_pattern("%")
        feat.add_pattern("本次.*?数量")
        feat.add_pattern("\d+\.\d+")    
        feat.add_pattern("公司股份")
        feat.add_pattern("累计")
        return feat
        
    def build_type_feature(self, tags):
        """构造实体类型特征
        这边的tags前期处理已经标记成日期、百分比等，这边只是简单做个判断

        Args
        ----
        tags : list
            取自特征数据矩阵的tag列，tags里每一个值代表该数据矩阵中某行的实体的类型

        Returns
        -------
        feature : np.array
            特征矩阵
        """
        entity_tag = ["date", "percent","stockquantity", "org", "amount"]
        feats = []
        for tag in tags:
            f = [0 for _ in entity_tag]
            f[entity_tag.index(tag)] = 1
            feats.append(f)
        return np.array(feats)
    
    def build_pattern_feature(self, left, right):
        """构造模式特征
        左邻和右邻分别构建

        Args
        ----
        left : str
            实体左边的句子

        right : str
            实体右边的句子

        Returns
        -------
        pfeat_left : np.array
            特征矩阵
        pfeat_right : np.array
            特征矩阵
        """
        pfeat_left = self.pattern_feat.build_feat_matrix(left)
        pfeat_right = self.pattern_feat.build_feat_matrix(right)
        return pfeat_left, pfeat_right

    def build_words_feature(self, left, right):
        """构造词语特征
        左邻和右邻分别构建, 然后合并做特征选择

        Args
        ----
        left : str
            实体左边的句子

        right : str
            实体右边的句子

        Returns
        -------
        feature : np.array
            特征矩阵
        """
        left = process_field_docs(left)
        right = process_field_docs(right)
        wordfeat_left = self.left_feat.transform(left)
        wordfeat_right = self.right_feat.transform(right)
        
        tmp_wordfeat = hstack((wordfeat_left, wordfeat_right))
        feature = self.feat_select.transform(tmp_wordfeat)
        return feature.toarray()
    
    def build_feature(self, left, right, entities_tag):
        """构造特征，组合了词语特征，模式特征，类型特征

        Args
        ----
        left : str
            实体左边的句子

        right : str
            实体右边的句子

        entities_tag : list
            取自特征数据矩阵的tag列，tags里每一个值代表该数据矩阵中某行的实体的类型

        Returns
        -------
        feature : np.array
            特征矩阵
        """
        typefeat = self.build_type_feature(entities_tag)
        pfeat_left, pfeat_right = self.build_pattern_feature(left, right)
        wordsfeat = self.build_words_feature(left, right)
        feature = np.concatenate((typefeat, pfeat_left, pfeat_right, wordsfeat), axis=1)
        return feature

    @classmethod
    def from_modelfile(cls, leftvectorizer_path=None, rightvectorizer_path=None, featselect_path=None):
        """快速创建实例的类方法
        """
        if leftvectorizer_path is None:
            leftvectorizer_path = "models/zengjianchi/text_field_leftfeat.m"
        if rightvectorizer_path is None:
            rightvectorizer_path = "models/zengjianchi/text_field_rightfeat.m"
        if featselect_path is None:
            featselect_path = "models/zengjianchi/text_field_featselect.m"
            
        leftvectorizer = joblib.load(leftvectorizer_path)
        rightvectorizer = joblib.load(rightvectorizer_path)
        select = joblib.load(featselect_path)
        feat = cls(leftvectorizer, rightvectorizer, select)
        return feat


if __name__ == "__main__":
    path = os.path.abspath("../data/training_data/zengjianchi_text_field2.csv")
    
    df = pd.read_csv(path)
    
    entity_tag = df["entity_tag"].tolist()
    label = df['label'].tolist()
    
    left = df['left'].fillna("").tolist()
    right = df['right'].fillna("").tolist()
    

    """ 训练分词特征构造器 """   
    docs_left = process_field_docs(left)
    docs_right = process_field_docs(right)
    
    vectorizer_left = CountVectorizer(binary=True)
    vectorizer_left.fit_transform(docs_left)
    
    vectorizer_right = CountVectorizer(binary=True)
    vectorizer_right.fit_transform(docs_right)
    
    wordfeat_left = vectorizer_left.transform(docs_left)
    wordfeat_right = vectorizer_right.transform(docs_right)
    
    tmp_wordfeat = hstack((wordfeat_left, wordfeat_right))

    select = SelectKBest(chi2, k=60)
    select.fit_transform(tmp_wordfeat, label)
    
    """构造训练特征"""
    field_feat = FieldFeature(vectorizer_left, vectorizer_right, select)
    features = field_feat.build_feature(left, right, entity_tag)

    """训练分类器"""
    clf = svm.SVC()    # 支持向量机分类

    # 交叉验证看效果
    k_fold = KFold(n_splits=5)
    scores = cross_val_score(clf, features, label, cv=k_fold, n_jobs=-1)
    print(scores)
    
    clf.fit(features, label)

    """模型保存"""
    joblib.dump(vectorizer_left, "models/zengjianchi/text_field_leftfeat.m")
    joblib.dump(vectorizer_right, "models/zengjianchi/text_field_rightfeat.m")
    joblib.dump(select, "models/zengjianchi/text_field_featselect.m")
    
    joblib.dump(clf, "models/zengjianchi/text_field_clf.m")
    
